import numpy as np
import pandas as pd
from scipy import optimize

def get_gerber_matrix(returns, threshold):
    std_returns = returns.std()
    thresholds = std_returns * threshold
    stock_num = returns.shape[1]
    gerber_matrix = pd.DataFrame(index=returns.columns, columns=returns.columns)
    for stock1 in returns.columns:
        for stock2 in returns.columns:
            uu, dd, ud, du, nn = 0, 0, 0, 0, 0
            for date in returns.index:                
                if returns.loc[date, stock1] > thresholds[stock1] and returns.loc[date, stock2] > thresholds[stock2]:
                    uu += 1
                elif returns.loc[date, stock1] < -thresholds[stock1] and returns.loc[date, stock2] < -thresholds[stock2]:
                    dd += 1
                elif returns.loc[date, stock1] > thresholds[stock1] and returns.loc[date, stock2] < -thresholds[stock2]:
                    ud += 1
                elif returns.loc[date, stock1] < -thresholds[stock1] and returns.loc[date, stock2] > thresholds[stock2]:
                    du += 1
                elif -thresholds[stock1] < returns.loc[date, stock1] < thresholds[stock1] and \
                     -thresholds[stock2] < returns.loc[date, stock2] < thresholds[stock2]:
                    nn += 1
            gerber_matrix.loc[stock1, stock2] = (uu+dd-ud-du)/(returns.shape[0]-nn)*std_returns[stock1]*std_returns[stock2]
    return gerber_matrix
            

def get_gerber_matrix_modified(returns, threshold, lambda1, lambda2):
    std_returns = returns.std()
    thresholds = std_returns * threshold
    stock_num = returns.shape[1]
    gerber_matrix = pd.DataFrame(index=returns.columns, columns=returns.columns)
    for stock1 in returns.columns:
        for stock2 in returns.columns:
            uu, dd, nn, un, dn, nd, nu, ud, du = 0, 0, 0, 0, 0, 0, 0, 0, 0
            for date in returns.index:                
                if returns.loc[date, stock1] > thresholds[stock1] and returns.loc[date, stock2] > thresholds[stock2]:
                    uu += 1
                elif returns.loc[date, stock1] < -thresholds[stock1] and returns.loc[date, stock2] < -thresholds[stock2]:
                    dd += 1
                elif -thresholds[stock1] < returns.loc[date, stock1] < thresholds[stock1] and \
                     -thresholds[stock2] < returns.loc[date, stock2] < thresholds[stock2]:
                    nn += 1
                elif returns.loc[date, stock1] > thresholds[stock1] and \
                     -thresholds[stock2] < returns.loc[date, stock2] < thresholds[stock2]:
                    un += 1
                elif returns.loc[date, stock1] < -thresholds[stock1] and \
                     -thresholds[stock2] < returns.loc[date, stock2] < thresholds[stock2]:
                    dn += 1
                elif -thresholds[stock1] < returns.loc[date, stock1] < thresholds[stock1] and \
                     returns.loc[date, stock2] > thresholds[stock2]:
                    nu += 1
                elif -thresholds[stock1] < returns.loc[date, stock1] < thresholds[stock1] and \
                     returns.loc[date, stock2] < -thresholds[stock2]:
                    nd += 1
                elif returns.loc[date, stock1] > thresholds[stock1] and returns.loc[date, stock2] < -thresholds[stock2]:
                    ud += 1
                elif returns.loc[date, stock1] < -thresholds[stock1] and returns.loc[date, stock2] > thresholds[stock2]:
                    du += 1                
            gerber_matrix.loc[stock1, stock2] = (uu+dd+lambda1*nn-lambda2*(nu+nd+un+dn)-ud-du)/(uu+dd+lambda1*nn+lambda2*(nu+nd+un+dn)+ud+du)*std_returns[stock1]*std_returns[stock2]
    return gerber_matrix


























































